{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8e3f10a3-fca1-4e2a-b39e-7cd8b44c724a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jupyter environment detected. Enabling Open3D WebVisualizer.\n",
      "[Open3D INFO] WebRTC GUI backend enabled.\n",
      "[Open3D INFO] WebRTCWindowSystem: HTTP handshake server disabled.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import open3d as o3d\n",
    "import os\n",
    "import time\n",
    "import json\n",
    "import cv2\n",
    "from tqdm import tqdm\n",
    "import math\n",
    "from scipy.spatial.transform import Rotation as R\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy\n",
    "import shutil\n",
    "\n",
    "from plyfile import PlyData, PlyElement\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3db0619f-a9bb-4aea-890d-d8c8cb3f7e18",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_intrinsic(data_folder):\n",
    "    with open(os.path.join(data_folder,'data/intrinsic/intrinsic_depth.txt'), 'r') as f:\n",
    "        lines = f.readlines()\n",
    "    intrinsic = np.zeros((4,4))\n",
    "    for i, line in enumerate(lines):\n",
    "        for j, content in enumerate(line.split(' ')):\n",
    "            intrinsic[i][j] = float(content)\n",
    "    return intrinsic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9574fb19-fe0c-45ae-9e59-3125a17e3566",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████| 5578/5578 [01:33<00:00, 59.58it/s]\n",
      "100%|███████████████████████████████████████| 1807/1807 [00:29<00:00, 61.95it/s]\n",
      "100%|███████████████████████████████████████| 2324/2324 [00:35<00:00, 65.18it/s]\n",
      "100%|███████████████████████████████████████| 2034/2034 [00:31<00:00, 63.85it/s]\n",
      "100%|███████████████████████████████████████| 2349/2349 [00:35<00:00, 65.97it/s]\n",
      "100%|███████████████████████████████████████| 1988/1988 [00:31<00:00, 63.08it/s]\n",
      "100%|███████████████████████████████████████| 7643/7643 [01:52<00:00, 68.11it/s]\n",
      "100%|███████████████████████████████████████| 6306/6306 [01:31<00:00, 69.10it/s]\n"
     ]
    }
   ],
   "source": [
    "fps_fake = 20 # fake camera frequency for offline ORB-SLAM\n",
    "raw_folder = \"./scans\"\n",
    "processed_folder = \"./processed\"\n",
    "\n",
    "scenes = ['scene0000_00', 'scene0054_00', 'scene0059_00', 'scene0106_00', 'scene0169_00', 'scene0181_00', 'scene0207_00', 'scene0233_00']\n",
    "\n",
    "for scene_idx, scene in enumerate(scenes):\n",
    "    save_folder = os.path.join(processed_folder,scene)\n",
    "    data_folder = os.path.join(raw_folder,scene)\n",
    "    \n",
    "    os.makedirs(save_folder)\n",
    "    os.makedirs(os.path.join(save_folder,\"rgb\"))\n",
    "    os.makedirs(os.path.join(save_folder,\"depth\"))\n",
    "    \n",
    "    shutil.copy(os.path.join(data_folder, 'data/intrinsic/intrinsic_depth.txt'), \n",
    "                os.path.join(save_folder, 'intrinsic.txt'))\n",
    "    \n",
    "    with open(os.path.join(save_folder,'gt_pose.txt'), 'w') as f:\n",
    "        f.write('# timestamp tx ty tz qx qy qz qw\\n')\n",
    "    \n",
    "    initial_time_stamp = time.time()  \n",
    "    \n",
    "    color_folder = os.path.join(data_folder,\"data/color\")\n",
    "    depth_folder = os.path.join(data_folder,\"data/depth\")\n",
    "    pose_folder = os.path.join(data_folder,\"data/pose\")\n",
    "    \n",
    "    num_frames = len(os.listdir(color_folder))\n",
    "    \n",
    "    frame_idx = 0\n",
    "    for raw_idx in tqdm(range(num_frames)):\n",
    "        with open(os.path.join(pose_folder,\"{}.txt\".format(raw_idx)), \"r\") as f:\n",
    "            lines = f.readlines()\n",
    "            M_w_c = np.zeros((4,4))\n",
    "            for i in range(4):\n",
    "                content = lines[i].split(\" \")\n",
    "                for j in range(4):\n",
    "                    M_w_c[i,j] = float(content[j])\n",
    "                    \n",
    "        if \"inf\" in lines[0]:\n",
    "            # invalid gt poses, skip this frame\n",
    "            continue\n",
    "\n",
    "        ######## convert depth to [m] and float type #########\n",
    "        depth = cv2.imread(os.path.join(depth_folder,\"{}.png\".format(raw_idx)),cv2.IMREAD_UNCHANGED)\n",
    "        depth = depth.astype(\"float32\")/1000.0\n",
    "\n",
    "        ######## resize rgb to the same size of depth #########\n",
    "        rgb = cv2.imread(os.path.join(color_folder,\"{}.jpg\".format(raw_idx)))\n",
    "        rgb = cv2.resize(rgb,(depth.shape[1],depth.shape[0]),interpolation=cv2.INTER_CUBIC)\n",
    "\n",
    "        cv2.imwrite(os.path.join(save_folder,\"rgb/frame_{}.png\".format(str(frame_idx).zfill(5))),rgb)\n",
    "        cv2.imwrite(os.path.join(save_folder,\"depth/frame_{}.TIFF\".format(str(frame_idx).zfill(5))),depth)\n",
    "\n",
    "        content = \"{:.4f}\".format(initial_time_stamp + frame_idx*1.0/fps_fake)\n",
    "        for t in M_w_c[:3,3]:\n",
    "            content += \" {:.9f}\".format(t)\n",
    "        for q in R.from_matrix(M_w_c[:3,:3]).as_quat():\n",
    "            content += \" {:.9f}\".format(q)\n",
    "        \n",
    "        with open(os.path.join(save_folder,'gt_pose.txt'), 'a') as f:\n",
    "            f.write(content + '\\n')\n",
    "            \n",
    "        frame_idx += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d1cc4cd-9868-4830-a849-d6e5702c5f20",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
